package util;

import model.Field;
import model.Figure;
import model.Position;
import org.apache.log4j.Logger;
import org.encog.engine.network.activation.ActivationSigmoid;
import org.encog.ml.MLRegression;
import org.encog.ml.data.basic.BasicMLData;
import org.encog.ml.data.buffer.BufferedMLDataSet;
import org.encog.ml.train.MLTrain;
import org.encog.neural.networks.BasicNetwork;
import org.encog.neural.networks.layers.BasicLayer;
import org.encog.neural.networks.training.propagation.resilient.ResilientPropagation;

import java.io.*;
import java.util.*;

/**
 * Class with static utilities methods
 *
 * @author Alex Vikharev alex.vikharev@gmail.com
 *         created 04.02.12
 */
public class Util {
    private static final Logger log = Logger.getLogger(Util.class);

    public static final BasicMLData IDEAL_WIN_DATA = new BasicMLData(new double[]{1.0});
    public static final BasicMLData IDEAL_LOSE_DATA = new BasicMLData(new double[]{0.0});
    public static final int[][] PERMUTATIONS = new int[24][4];

    static {
        PERMUTATIONS[0] = new int[]{0, 1, 2, 3};
        PERMUTATIONS[1] = new int[]{0, 1, 3, 2};
        PERMUTATIONS[2] = new int[]{0, 2, 1, 3};
        PERMUTATIONS[3] = new int[]{0, 2, 3, 1};
        PERMUTATIONS[4] = new int[]{0, 3, 1, 2};
        PERMUTATIONS[5] = new int[]{0, 3, 2, 1};
        PERMUTATIONS[6] = new int[]{1, 0, 2, 3};
        PERMUTATIONS[7] = new int[]{1, 0, 3, 2};
        PERMUTATIONS[8] = new int[]{1, 2, 0, 3};
        PERMUTATIONS[9] = new int[]{1, 2, 3, 0};
        PERMUTATIONS[10] = new int[]{1, 3, 0, 2};
        PERMUTATIONS[11] = new int[]{1, 3, 2, 0};
        PERMUTATIONS[12] = new int[]{2, 0, 1, 3};
        PERMUTATIONS[13] = new int[]{2, 0, 3, 1};
        PERMUTATIONS[14] = new int[]{2, 1, 0, 3};
        PERMUTATIONS[15] = new int[]{2, 1, 3, 0};
        PERMUTATIONS[16] = new int[]{2, 3, 0, 1};
        PERMUTATIONS[17] = new int[]{2, 3, 1, 0};
        PERMUTATIONS[18] = new int[]{3, 0, 1, 2};
        PERMUTATIONS[19] = new int[]{3, 0, 2, 1};
        PERMUTATIONS[20] = new int[]{3, 1, 0, 2};
        PERMUTATIONS[21] = new int[]{3, 1, 2, 0};
        PERMUTATIONS[22] = new int[]{3, 2, 0, 1};
        PERMUTATIONS[23] = new int[]{3, 2, 1, 0};
    }

    public static Position getCanonical(Position move) {
        List<Field> fields = move.getField().getAllSymmetric();
        SortedSet<Position> set = new TreeSet<>(Position.comparator());

        for (Field field : fields) {
            for (int xor = 0; xor < 16; xor++) {
                Field newField = field.xorBitsInFigures(xor);
                Figure figure = move.getFigure() == null ? null : move.getFigure().xor(xor);
                for (int[] permutation : PERMUTATIONS) {
                    Field permutedField = newField.permuteBitsInFigures(permutation);
                    Figure permutedFigure = figure == null ? null : figure.permute(permutation);
                    set.add(new Position(permutedField, permutedFigure));
                }
            }
        }

        return set.first();
    }

    public static List<Field> permutedBits(Field field) {
        List<Field> result = new ArrayList<>();
        for (int[] permutation : PERMUTATIONS) {
            result.add(field.permuteBitsInFigures(permutation));
        }
        return result;
    }

    public static boolean hasCommonPart(Figure figure1, Figure figure2, Figure figure3, Figure figure4) {
        return (figure1 != null && figure2 != null && figure3 != null && figure4 != null) &&
                ((figure1.code & figure2.code & figure3.code & figure4.code) > 0 ||
                        (figure1.code | figure2.code | figure3.code | figure4.code) < 0b1111);
    }

    public static List<Figure[]> getWinFigures() {
        List<Figure[]> result = new ArrayList<>(536);
        for (int f1 = 0; f1 < 13; f1++) {
            for (int f2 = f1 + 1; f2 < 14; f2++) {
                for (int f3 = f2 + 1; f3 < 15; f3++) {
                    for (int f4 = f3 + 1; f4 < 16; f4++) {
                        Figure[] a = new Figure[]{Figure.byCode(f1), Figure.byCode(f2), Figure.byCode(f3), Figure.byCode(f4)};
                        if (hasCommonPart(a[0], a[1], a[2], a[3])) {
                            result.add(a);
                        }
                    }
                }
            }
        }
        return result;
    }

    public static void makeTrainDataFile(File file) {
        Set<Position> moves = new HashSet<>();

        BufferedMLDataSet dataSet = new BufferedMLDataSet(file);
        dataSet.beginLoad(17, 1);

        List<Figure[]> winFigures = getWinFigures();
        for (Figure[] figures : winFigures) {
            List<Figure[]> permutations = permutations(figures);
            for (Figure[] permutation : permutations) {

                for (int y = 0; y < 4; y++) {
                    Field.FieldBuilder builder = Field.builder();
                    for (int x = 0; x < 4; x++) {
                        builder.setFigure(x, y, permutation[x]);
                    }
                    Field field = builder.build();
                    addTrainData(dataSet, field, moves, BITWISE_CONVERTER);
                }

                for (int x = 0; x < 4; x++) {
                    Field.FieldBuilder builder = Field.builder();
                    for (int y = 0; y < 4; y++) {
                        builder.setFigure(x, y, permutation[y]);
                    }
                    Field field = builder.build();
                    addTrainData(dataSet, field, moves, BITWISE_CONVERTER);
                }

                Field.FieldBuilder diag1Builder = Field.builder();
                Field.FieldBuilder diag2Builder = Field.builder();

                for (int i = 0; i < 4; i++) {
                    diag1Builder.setFigure(i, i, permutation[i]);
                    diag2Builder.setFigure(i, 3 - i, permutation[i]);
                }

                addTrainData(dataSet, diag1Builder.build(), moves, BITWISE_CONVERTER);
                addTrainData(dataSet, diag2Builder.build(), moves, BITWISE_CONVERTER);

            }
        }
        dataSet.endLoad();
        log.info("Created file with: " + dataSet.getRecordCount() + " training samples");
    }

    private static void addTrainData(BufferedMLDataSet dataSet, Field field, Set<Position> usedMoves, Function<Position, double[]> converter) {
        List<Position> moves = makeMove(field);
        for (Position move : moves) {
            Position canonical = getCanonical(move);
            if (usedMoves.add(canonical)) {
                dataSet.add(new BasicMLData(converter.apply(canonical)), IDEAL_WIN_DATA);
            }
        }
        for (int x = 0; x < 4; x++) {
            for (int y = 0; y < 4; y++) {
                if (field.get(x, y) != null) {
                    Figure figure = field.get(x, y);
                    Field newField = field.removeFigure(x, y);
                    Position move = new Position(newField, figure);
                    if (usedMoves.add(move)) {
                        dataSet.add(new BasicMLData(converter.apply(move)), IDEAL_LOSE_DATA);
                    }
                }
            }
        }

    }

    private static List<Position> makeMove(Field field) {
        List<Position> moves = new ArrayList<>();
        EnumSet<Figure> freeFigures = EnumSet.complementOf(EnumSet.copyOf(field.getUsedFigures()));
        for (Figure usedFigure : freeFigures) {
            moves.add(new Position(field, usedFigure));
        }
        return moves;
    }

    public static List<Figure[]> permutations(Figure[] figures) {
        List<Figure[]> result = new ArrayList<>(24);
        for (int[] permutation : PERMUTATIONS) {
            result.add(new Figure[]{figures[permutation[0]], figures[permutation[1]], figures[permutation[2]], figures[permutation[3]]});
        }
        return result;
    }

    public static void main(String[] args) throws FileNotFoundException {
//        makeTrainDataFile(new File("train2.enc"));

        teachNetwork(new File("train2.enc"), "17-17-17-1_400_iter.weights");
    }

    private static void teachNetwork(File file, String outputFileName) throws FileNotFoundException {
        BasicNetwork network = new BasicNetwork();
        network.addLayer(new BasicLayer(17));
        network.addLayer(new BasicLayer(new ActivationSigmoid(), true, 17));
        network.addLayer(new BasicLayer(new ActivationSigmoid(), true, 17));
        network.addLayer(new BasicLayer(new ActivationSigmoid(), true, 1));
        network.getStructure().finalizeStructure();
        network.reset();

        BufferedMLDataSet dataSet = new BufferedMLDataSet(file);
        int epoch = 0;
        MLTrain train = new ResilientPropagation(network, dataSet);
//        MLTrain train = new Backpropagation(network, dataSet);
        log.info("Start teaching bot");
        log.info("Training network, epoch: " + epoch + " error: " + train.getError());
        for (int i = 0; i < 50; i++) {
            train.iteration();
            ++epoch;
            log.info("Training network, epoch: " + epoch + " error: " + train.getError());

        }
        log.info("Training is over, error: " + train.getError());

        saveClassifier(network, outputFileName);
    }

    public static final Function<Position, Position> TO_CANONICAL = new Function<Position, Position>() {
        @Override
        public Position apply(Position arg) {
            return Util.getCanonical(arg);
        }
    };

    public static final Function<Position, double[]> BITWISE_CONVERTER = new Function<Position, double[]>() {

        @Override
        public double[] apply(Position move) {
            double[] result = new double[17 * 5];
            Field field = move.getField();
            int i = 0;
            for (int x = 0; x < 4; x++) {
                for (int y = 0; y < 4; y++) {
                    Figure figure = field.get(x, y);
                    if (figure == null) {
                        i += 5;
                    } else {
                        result[i++] = 1.0;
                        for (double bit : figure.bits) {
                            result[i++] = bit;
                        }
                    }
                }
            }
            if (move.getFigure() != null) {
                result[i++] = 1.0;
                for (double bit : move.getFigure().bits) {
                    result[i++] = bit;
                }
            }
            return result;
        }
    };


    public static final Function<Position, double[]> SMOOTH_CONVERTER = new Function<Position, double[]>() {
        @Override
        public double[] apply(Position move) {
            double[] result = new double[17];
            int i = 0;
            for (int x = 0; x < 4; x++) {
                for (int y = 0; y < 4; y++) {
                    Figure figure = move.getField().get(x, y);
                    double value = figure == null ? -1 : figure.code;
                    result[i++] = value;
                }
            }
            Figure figure = move.getFigure();
            double value = figure == null ? -1 : figure.code;
            result[i] = value;
            return result;
        }
    };

    public static void saveClassifier(MLRegression classifier, String fileName) {
        try {
            FileOutputStream fout = new FileOutputStream(fileName);
            ObjectOutputStream oos = new ObjectOutputStream(fout);
            oos.writeObject(classifier);
            oos.close();
        } catch (Exception e) {
            log.error(e);
        }
    }

    public static MLRegression loadClassifier(String fileName) {
        try {
            FileInputStream fin = new FileInputStream(fileName);
            ObjectInputStream ois = new ObjectInputStream(fin);
            MLRegression classifier = (MLRegression) ois.readObject();
            ois.close();
            return classifier;
        } catch (Exception e) {
            log.error(e);
            return null;
        }
    }
}
